import numpy as np
from math import cos, pi, sqrt


def dct_2d(img):
    height = img.shape[0]
    width = img.shape[1]
    imageRow = np.zeros_like(img).astype(float)
    imageCol = np.zeros_like(img).astype(float)

    for h in range(height):
        imageRow[h, :] = dct_1d(img[h, :])
    for w in range(width):
        imageCol[:, w] = dct_1d(imageRow[:, w])

    return imageCol


def dct_1d(img):
    n = len(img)
    retImg = np.zeros_like(img).astype(float)

    for k in range(n):
        sum = 0
        for i in range(n):
            sum += img[i] * cos(2 * pi * k / (2.0 * n) * i + (k * pi) / (2.0 * n))
        ck = sqrt(0.5) if k == 0 else 1
        retImg[k] = sqrt(2.0 / n) * ck * sum

    return retImg
